"""
We need the last function to help extract the final answer of SPARQL, used in check_sparql
"""

import argparse
import json
import os
import pickle
from itertools import chain

import numpy as np
from tqdm import tqdm
from transformers import BartTokenizerFast


def get_program_seq(program):
    seq = []
    for item in program:
        func = item["function"]
        inputs = item["inputs"]
        seq.append(func + "(" + "<c>".join(inputs) + ")")
    seq = "<b>".join(seq)
    # print(program)
    # print(seq)
    return seq

valid_ids = json.load(
    open("../LLM_KGQA/save/kqapro-valid/9types-gpt-3.5-turbo-16k-v2.0/train_ids.json")
)
valid_ids += json.load(
    open("../LLM_KGQA/save/kqapro-valid/9types-gpt-3.5-turbo-16k-v2.0/dev_ids.json")
)
valid_ids = set(valid_ids)


def encode_dataset(dataset, vocab, tokenizer: BartTokenizerFast, test=False):
    # dataset = [d for d in dataset if d["id"] in valid_ids]
    questions = []
    programs = []
    for item in tqdm(dataset):
        question = item["question"]
        questions.append(question)
        if not test:
            program = item["program"]
            program = get_program_seq(program)
            programs.append(program)

    _questions = tokenizer(questions, padding=True)
    max_question_length = len(_questions["input_ids"][0])
    _programs = tokenizer(programs, padding=True)
    max_programs_length = len(_programs["input_ids"][0])

    # 上面只是为了得到最大长度，下面才是真正的编码
    questions = []
    programs = []
    choices = []
    answers = []
    for item in tqdm(dataset):
        question = item["question"]
        questions.append(question)
        _ = [vocab["answer_token_to_idx"][w] for w in item["choices"]]
        choices.append(_)
        if not test:
            program = item["program"]
            program = get_program_seq(program)
            programs.append(program)
            answers.append(vocab["answer_token_to_idx"].get(item["answer"]))

    input_ids = tokenizer.batch_encode_plus(
        questions, max_length=max_question_length, padding=True, truncation=True
    )
    source_ids = np.array(input_ids["input_ids"], dtype=np.int32)
    source_mask = np.array(input_ids["attention_mask"], dtype=np.int32)
    if not test:
        target_ids = tokenizer.batch_encode_plus(
            programs, max_length=max_programs_length, padding=True, truncation=True
        )
        target_ids = np.array(target_ids["input_ids"], dtype=np.int32)
    else:
        target_ids = np.array([], dtype=np.int32)
    choices = np.array(choices, dtype=np.int32)
    answers = np.array(answers, dtype=np.int32)
    return source_ids, source_mask, target_ids, choices, answers


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dir", required=True)
    parser.add_argument("--output_dir", required=True)
    parser.add_argument("--model_name_or_path", required=True)
    args = parser.parse_args()

    print("Build kb vocabulary")
    vocab = {"answer_token_to_idx": {}}
    print("Load questions")
    train_set = json.load(open(os.path.join(args.input_dir, "train.json")))
    val_set = json.load(open(os.path.join(args.input_dir, "val.json")))
    test_set = json.load(open(os.path.join(args.input_dir, "test.json")))
    for question in chain(train_set, val_set, test_set):
        for a in question["choices"]:
            if not a in vocab["answer_token_to_idx"]:
                vocab["answer_token_to_idx"][a] = len(vocab["answer_token_to_idx"])

    if not os.path.isdir(args.output_dir):
        os.mkdir(args.output_dir)
    fn = os.path.join(args.output_dir, "vocab.json")
    print("Dump vocab to {}".format(fn))
    with open(fn, "w") as f:
        json.dump(vocab, f, indent=2)
    for k in vocab:
        print("{}:{}".format(k, len(vocab[k])))
    tokenizer = BartTokenizerFast.from_pretrained(args.model_name_or_path)
    for name, dataset in zip(("train", "val", "test"), (train_set, val_set, test_set)):
        
        # myadd
        if name == "test":
            continue

        print("Encode {} set".format(name))
        outputs = encode_dataset(dataset, vocab, tokenizer, name == "test")
        assert len(outputs) == 5
        print(
            "shape of input_ids of questions, attention_mask of questions, input_ids of sparqls, choices and answers:"
        )
        with open(os.path.join(args.output_dir, "{}.pt".format(name)), "wb") as f:
            for o in outputs:
                print(o.shape)
                pickle.dump(o, f)


if __name__ == "__main__":
    main()
